import numpy as np
from pyDOE import lhs

# 假设有41种产品，编号从1到41
num_products = 41
product_numbers = list(range(1, num_products + 1))  # 产品编号的列表

# 定义年份范围和训练样本数量
years = list(range(2024, 2031))  # 从2024到2030年
n_train_samples = 1  # 样本数量，可以根据需要调整
n_variables = 4  # 变量数量（价格变化、销售量变化、亩产量变化、种植成本变化）

# 使用拉丁超立方采样生成样本
lhs_samples_train = lhs(
    n_variables * num_products * len(years), samples=n_train_samples
)

# 重塑采样结果为（样本数，产品数，年份，变量数）的形状
lhs_samples_train = lhs_samples_train.reshape(
    n_train_samples, num_products, len(years), n_variables
)

# 定义敏感性系数字典
price_cost_sensitivity = {}
sales_price_sensitivity = {}

# 为每个产品设置随机种子并生成敏感性系数
for product_num in product_numbers:
    np.random.seed(product_num)  # 使用产品编号作为随机种子
    price_cost_sensitivity[product_num] = np.random.uniform(
        0.03, 0.07
    )  # 售价随成本变化的系数
    sales_price_sensitivity[product_num] = np.random.uniform(
        -0.15, -0.05
    )  # 销量随售价变化的系数

# 定义交叉弹性矩阵（每对产品的交叉弹性系数）
cross_elasticity = np.random.uniform(
    -0.2, 0.2, (num_products, num_products)
)  # 交叉弹性范围从-0.2到0.2

# 初始化样本列表
cumulative_samples_list_train = []

for i in range(n_train_samples):
    # 初始化每个样本的数据字典
    year_data = {year: {} for year in years}

    # 初始化每个产品的累积变化
    cumulative_price = {product: 1.0 for product in product_numbers}
    cumulative_sales = {product: 1.0 for product in product_numbers}
    cumulative_yield = {product: 1.0 for product in product_numbers}
    cumulative_cost = {product: 1.0 for product in product_numbers}

    for year_idx, year in enumerate(years):
        # 记录每个产品上一次的价格和销量
        last_price = cumulative_price.copy()
        last_sales = cumulative_sales.copy()
        # 第一步：更新每个产品的价格和销量变化（不考虑交叉弹性影响）
        for product_idx, product_num in enumerate(product_numbers):
            # 生成对应产品和年份的采样数据
            price_change = 0.01 + lhs_samples_train[i, product_idx, year_idx, 0] * 0.04
            sales_change = -0.05 + lhs_samples_train[i, product_idx, year_idx, 1] * 0.10
            yield_change = -0.10 + lhs_samples_train[i, product_idx, year_idx, 2] * 0.20
            cost_change = 1.05

            # 更新种植成本的累积变化
            cumulative_cost[product_num] *= cost_change

            # 判断并应用各类作物的变化率
            if product_num in [6, 7]:  # 小麦和玉米（编号为6和7的产品）
                cumulative_price[product_num] *= 1  # 价格基本稳定
                cumulative_sales[product_num] *= (
                    1.05 + lhs_samples_train[i, product_idx, year_idx, 1] * 0.05
                )  # 销量增加5%~10%
                cumulative_yield[product_num] *= 1 + yield_change  # 亩产量变化±10%

            elif product_num in range(17, 38):  # 蔬菜类作物（编号为17到37的产品）
                cumulative_price[product_num] *= 1.05  # 价格增加5%
                cumulative_sales[product_num] *= 1 + sales_change  # 销量变化±5%
                cumulative_yield[product_num] *= 1 + yield_change  # 亩产量变化±10%

            elif product_num in range(38, 42):  # 食用菌类作物（编号为38到41的产品）
                cumulative_price[product_num] *= (
                    0.99 - lhs_samples_train[i, product_idx, year_idx, 0] * 0.05
                )  # 价格下降1%-6%
                cumulative_sales[product_num] *= 1 + sales_change  # 销量变化±5%
                cumulative_yield[product_num] *= 1 + yield_change  # 亩产量变化±10%

            else:  # 其他未分类的作物
                cumulative_price[product_num] *= (
                    0.95 + lhs_samples_train[i, product_idx, year_idx, 0] * 0.10
                )
                cumulative_sales[product_num] *= 1 + sales_change  # 销量变化±5%
                cumulative_yield[product_num] *= 1 + yield_change  # 亩产量变化±10%

            # 根据成本变化影响售价
            cumulative_price[product_num] *= 1 + (
                price_cost_sensitivity[product_num] * (cumulative_cost[product_num] - 1)
            )

            # 根据售价变化影响销量
            cumulative_sales[product_num] *= (
                1
                + (cumulative_price[product_num] / last_price[product_num] - 1)
                * sales_price_sensitivity[product_num]
            )

        # 第二步：根据交叉弹性更新销量
        for product_num in product_numbers:
            for other_product_num in product_numbers:
                if other_product_num != product_num:  # 不考虑自身对自身的交叉弹性
                    cross_effect = cross_elasticity[
                        product_num - 1, other_product_num - 1
                    ]
                    other_price_change = (
                        cumulative_price[other_product_num]
                        / last_price[other_product_num]
                        - 1
                    )
                    cumulative_sales[product_num] *= (
                        1 + cross_effect * other_price_change
                    )

            # 记录每年的数据
            year_data[year][product_num] = {
                "价格变化": round(cumulative_price[product_num], 4),
                "销售量变化": round(cumulative_sales[product_num], 4),
                "亩产量变化": round(cumulative_yield[product_num], 4),
                "种植成本变化": round(cumulative_cost[product_num], 4),
            }

    # 将每个样本的结果加入列表
    cumulative_samples_list_train.append(year_data)

# 输出示例
for sample in cumulative_samples_list_train:
    print(sample)
